#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Author  : jeffzhang
# @Time    : 2018/04/03
# @File    : vulnerability_management.py
# @Desc    : ""

from flask import Blueprint, render_template, request, redirect, url_for, jsonify
from bson import ObjectId
from lib.mongo_db import connectiondb, db_name_conf
from InsectsAwake.views.authenticate import login_check

vulnerability_management = Blueprint('vulnerability_management', __name__)
vul_db = db_name_conf()['vul_db']


@vulnerability_management.route('/vulnerability-management', methods=['POST', 'GET'])
@login_check
def vul_list():
    # 删除漏洞
    if request.method == "GET":
        if request.args.get('delete'):
            vul_id = request.args.get('delete')
            task_id = connectiondb(vul_db).find_one({'_id': ObjectId(vul_id)})['task_id']
            connectiondb(vul_db).delete_one({'_id': ObjectId(vul_id)})
            return redirect(url_for('vulnerability_management.vul_list', task=task_id))

        # 漏洞复测
        elif request.args.get('rescan'):
            vul_id = request.args.get('rescan')
            print(vul_id)

        # 获取扫描结果详情
        elif request.args.get('result'):
            vul_id = request.args.get('result')
            vul_data = connectiondb(vul_db).find_one({'_id': ObjectId(vul_id)})
            del vul_data['task_id']
            del vul_data['_id']
            return jsonify(vul_data)

        # 从任务界面跳转过来,筛选出某任务中发现的漏洞
        elif request.args.get('task'):
            task_id = request.args.get('task')
            vul_data = connectiondb(vul_db).find({'task_id': ObjectId(task_id)}).sort('scan_date', -1)
            return render_template('vulnerability-management.html', vul_data=vul_data)

        # 从插件界面跳转过来,筛选出某个插件发现的漏洞
        elif request.args.get('plugin'):
            plugin_id = request.args.get('plugin')
            vul_data = connectiondb(vul_db).find({'plugin_id': plugin_id}).sort('scan_date', -1)
            return render_template('vulnerability-management.html', vul_data=vul_data)

        # 默认返回漏洞列表
        vul_data = connectiondb(vul_db).find().sort('scan_date', -1)
        return render_template('vulnerability-management.html', vul_data=vul_data)

    elif request.method == "POST":
        # 删除多选
        vul_id = request.form.get('vul_id').split(',', -1)
        try:
            for i in vul_id:
                connectiondb(vul_db).remove({'_id': ObjectId(i)})
        except Exception as e:
            print(e)
        return jsonify({'result': 'success'})
